import glob
import platform
import time
import joblib
import os
from sklearn.svm import SVC
import numpy as np
import variable as va

# 把结果写入到
def write_to_txt(list):
    with open('result.txt', 'w') as f:
        f.writelines(list)
    print('每张图片的识别结果存放在result.txt里面')

# 加载数据集特征文件
def load_feats(filePath):
    features = []
    labels = []
    # 加载所有的特征文件
    for feat_path in glob.glob(os.path.join(filePath, '*.feat')):
        data = joblib.load(feat_path)
        # 每个特征文件中都有一个样本，最后一个数据是label，其余的是特征
        features.append(data[:-1])
        labels.append(data[-1])
    # 返回特征与标签的列表
    return features, labels 
    

# 训练
def train():
    print('----- 开始训练 SVM 分类器 -----')
    t0 = time.time()
    # 加载训练集的特征与标签
    features, labels = load_feats(va.train_feat_path)
    
    # 借鉴https://blog.csdn.net/ustbclearwang/article/details/81236732参数使用
    # 定义分类器
    # 最佳kernel=rbf 最佳gamma=1.2 0.812500
    clf = SVC(C=1, kernel='rbf', gamma=1.2) # 0.81
    # clf = SVC(C=1, kernel='linear', gamma=1) # 0.59
    # clf = SVC(C=1, kernel='poly', gamma=1) # 0.53
    # clf = SVC(C=1, kernel='sigmoid', gamma=1) # 0.75
    # 对！训练只要一句
    clf.fit(features, labels)
    
    # 如果保存模型的文件夹不存在，创建这个文件夹
    if not os.path.exists(va.model_path):
        os.makedirs(va.model_path)
    # 保存模型
    joblib.dump(clf, va.model_path + 'model')
    t1 = time.time()
    print("训练之后的模型存放在model文件夹中")
    print('训练模型耗时 : %f 秒' % (t1 - t0))
    return clf


# 测试
def test(clf):
    t0 = time.time()
    correct_number = 0 # 分类正确的次数
    total = 0          # 分类总次数
    result_list = []   # 分类结果列表
    for feat_path in glob.glob(os.path.join(va.test_feat_path, '*.feat')):
        total += 1
        if platform.system() == 'Windows':
            symbol = '\\'
        else:
            symbol = '/'
        # 获取测试集中图像的特征
        image_name = feat_path.split(symbol)[1].split('.feat')[0]
        data_test = joblib.load(feat_path)

        # 规范化label
        data_test_feat = data_test[:-1].reshape((1, -1)).astype(np.float64)
        # 预测结果 -- 会得到一个浮点数
        result = clf.predict(data_test_feat)
        # 记录该图片的预测结果
        result_list.append(image_name + ' ' + va.labels[int(result[0])] + '\n')
        if int(result[0]) == data_test[-1]:
           correct_number += 1
    write_to_txt(result_list)
    rate = float(correct_number) / total
    t1 = time.time()
    print('准确率是： %f' % rate)
    print('测试耗时 : %f ' % (t1 - t0))
